from typing import Annotated, Literal, Optional
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage
from IPython.display import Image, display
import os

# 1. 定义状态类
class State(TypedDict):
    # 消息列表，使用 add_messages 函数来追加消息而不是覆盖
    messages: Annotated[list, add_messages]
    # 对话状态，用于决定流程走向
    conversation_stage: Literal["start", "ongoing", "ended"]

# 2. 定义节点函数

def user_input_node(state: State) -> State:
    """接收用户输入的节点"""
    if state["conversation_stage"] == "start":
        print("欢迎使用 LangGraph 对话演示! 输入 'exit' 结束对话。")
        state["conversation_stage"] = "ongoing"
        state["next"] = "user_input_node"
    user_input = input("用户: ")
    if user_input.lower() == 'exit':
        state["conversation_stage"] = "ended"
        return state

    # 将用户输入添加到消息列表
    state["messages"].append(HumanMessage(content=user_input))
    return state

def llm_node(state: State) -> State:
    """使用 LLM 生成回复的节点"""
    # 检查是否有API密钥
    api_key = os.getenv("OPENROUTER_API_KEY") or os.getenv("OPENAI_API_KEY")
    if api_key:
        # 如果有API密钥，使用真实的LLM
        try:
            llm = ChatOpenAI(
                model="moonshotai/kimi-k2:free",  # OpenRouter 模型格式: provider/model
                temperature=0.7,
                api_key=api_key,
                base_url="https://openrouter.ai/api/v1",  # OpenRouter API 基础 URL
            )
            response = llm.invoke(state["messages"])
            ai_response = response.content
        except Exception as e:
            print(f"LLM调用失败: {e}")
            ai_response = "抱歉，我现在无法正常回复。请检查API配置。"
    else:
        # 如果没有API密钥，使用模拟回复
        user_message = state["messages"][-1].content if state["messages"] else ""
        ai_response = f"你好！我收到了你的消息：'{user_message}'。这是一个模拟回复，因为没有配置API密钥。"
    
    # 将 LLM 回复添加到消息列表
    state["messages"].append(AIMessage(content=ai_response))
    print(f"AI: {ai_response}")
    return state

def decision_node(state: State) -> dict:
    """决定流程走向的节点"""
    if state["conversation_stage"] == "ended":
        print("对话已结束。")
        return {"next": END}
    return {"next": "user_input_node"}

# 3. 构建图

def build_graph() -> StateGraph:
    graph = StateGraph(State)

    # 添加节点
    graph.add_node("user_input_node", user_input_node)
    graph.add_node("llm_node", llm_node)
    graph.add_node("decision_node", decision_node)

    # 添加边
    graph.add_edge(START, "user_input_node")
    graph.add_edge("user_input_node", "llm_node")
    graph.add_edge("llm_node", "decision_node")
    graph.add_conditional_edges(
        "decision_node",
        lambda state: state["next"],
        {
            "user_input_node": "user_input_node",
            END: END,
        }
    )

    return graph

# 5. 可视化图结构

def visualize_graph(app, image_path: str = "graph_visualization.png") -> None:
    """生成并显示图结构图片

    Args:
        app: 编译后的图应用对象
        image_path: 图片保存路径
    """
    try:
        # 使用编译后的应用对象生成图片
        png_data = app.get_graph().draw_mermaid_png()
        
        # 保存图片到本地
        with open(image_path, "wb") as f:
            f.write(png_data)
        print(f"图结构已保存到 {image_path}")

        # 如果在 IPython 环境中，显示图片
        try:
            display(Image(image_path))
        except:
            print("无法在当前环境中显示图片，但已保存到文件")
    except Exception as e:
        print(f"生成图结构图片时出错: {e}")
        print("提示：可能需要安装 pygraphviz 或其他图形库")

def main():
    # 初始化图
    graph = build_graph()

    # 编译图
    app = graph.compile()
    
    # 可视化图结构
    # visualize_graph(app)

    # 运行图
    initial_state = {
        "messages": [],
        "conversation_stage": "start"
    }
    app.invoke(initial_state)

if __name__ == "__main__":
    main()
